package com.wengu.algorithm.alg;

import com.wengu.algorithm.dto.TaskAlg;
import com.wengu.algorithm.dto.UserAlg;
import com.wengu.algorithm.entity.AssignmentOutcome;
import com.wengu.algorithm.service.AssignmentOutcomeUserService;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.stream.Collectors;

/**
 * B-QEA算法实现
 * 结合合作蚁群搜索和基于执行成本权重任务密度的算法
 */
@Slf4j
public class BQEA extends QEA {
    // 当前处理的用户引用
    public UserAlg currentUser;


    // 全局信息素矩阵
    private Map<Long, Map<Long, Double>> globalPheromoneMatrix;
    // 局部信息素矩阵
    private Map<Long, Map<Long, Double>> localPheromoneMatrix;

    /**
     * 默认构造函数
     */
    public BQEA() {
        this(1, 10, 5, 0.1, 1.0, 2.0, 0.9, 0.1);
    }
    public BQEA(double alpha){
        this(alpha,30, 30, 0.1, 1.0, 2.0, 0.9, 0.1);
    }
    /**
     * 带参数的构造函数
     * @param alpha 权重调节参数
     * @param antCount 蚂蚁数量
     * @param roundCount 迭代轮数
     * @param initialPheromone 初始信息素值
     * @param pheromoneWeight 信息素权重
     * @param heuristicWeight 启发式权重(beta)
     * @param evaporationRate 信息素挥发率
     * @param bestSelectionProbability 最优路径选择概率
     */
    public BQEA(double alpha, int antCount, int roundCount,
                double initialPheromone, double pheromoneWeight,
                double heuristicWeight, double evaporationRate,
                double bestSelectionProbability) {
        super(alpha,bestSelectionProbability);
        this.antCount = antCount;
        this.roundCount = roundCount;
        this.initialPheromone = initialPheromone;
        this.pheromoneWeight = pheromoneWeight;
        this.heuristicWeight = heuristicWeight;
        this.evaporationRate = evaporationRate;
//        this.bestSelectionProbability = bestSelectionProbability;
    }

    @Override
    public List<TaskAlg> allocateTasks(UserAlg user, List<TaskAlg> allTasks,
                                       AssignmentOutcome assignmentOutcome,
                                       AssignmentOutcomeUserService assignmentOutcomeUserService,
                                       Map<Long, TaskAlg> taskCacheMap,
                                       String algName, Set<Long> notIncreaseTaskIdSet) {
        this.currentUser = user;

        // 初始化信息素矩阵
        initializePheromoneMatrix(allTasks);
        // 存储所有蚂蚁的路径上的任务的参数，包含执行次数和效用值
//        Map<String,Map<Long, TaskAlg>> antPathTaskParams = new HashMap<>();
        // 存储所有蚂蚁的路径
        List<List<TaskAlg>> allAntPaths = new ArrayList<>();
        List<Double> allAntPathUtilities = new ArrayList<>();
        spatialIndex = null;
        // 执行多轮迭代
        for (int round = 0; round < roundCount; round++) {
            // 每轮开始时，局部信息素更新为全局信息素
            copyGlobalToLocalPheromone();

            // 每只蚂蚁构建路径
            for (int ant = 0; ant < antCount; ant++) {
                List<TaskAlg> antPath = constructAntPath(user, allTasks, taskCacheMap, notIncreaseTaskIdSet);
                double pathUtility = calculateTotalDeltaF(antPath);

                allAntPaths.add(antPath);
                allAntPathUtilities.add(pathUtility);
            }

            // 更新全局信息素
            updateGlobalPheromone(allAntPaths, allAntPathUtilities);
        }

        // 选择最优路径
        List<TaskAlg> bestPath = selectBestPath(allAntPaths, allAntPathUtilities);

        // 计算用户总质量值
        double userTotalDeltaF = calculateTotalDeltaF(bestPath);

        // 更新任务执行次数和用户成本
        updateTaskAndUserStatus(user, bestPath, taskCacheMap, notIncreaseTaskIdSet);

        // 保存用户执行结果
        assignmentOutcomeUserService.saveOutcomeUser(user, assignmentOutcome, bestPath, userTotalDeltaF);
        log.info("用户[{}]执行平台安排的任务后的质量增量总和为[{}]",user.id,userTotalDeltaF);
        return bestPath;
    }

    /**
     * 初始化信息素矩阵
     * @param tasks 所有任务列表
     */
    private void initializePheromoneMatrix(List<TaskAlg> tasks) {
        globalPheromoneMatrix = new HashMap<>();
        localPheromoneMatrix = new HashMap<>();

        // 初始化所有任务间的信息素
        for (TaskAlg task1 : tasks) {
            Map<Long, Double> row1 = new HashMap<>();
            Map<Long, Double> row2 = new HashMap<>();

            for (TaskAlg task2 : tasks) {
                if (!task1.equals(task2)) {
                    row1.put(task2.id, initialPheromone);
                    row2.put(task2.id, initialPheromone);
                }
            }

            globalPheromoneMatrix.put(task1.id, row1);
            localPheromoneMatrix.put(task1.id, row2);
        }
    }

    /**
     * 将全局信息素复制到局部信息素
     */
    private void copyGlobalToLocalPheromone() {
        localPheromoneMatrix = new HashMap<>();

        for (Map.Entry<Long, Map<Long, Double>> entry : globalPheromoneMatrix.entrySet()) {
            Map<Long, Double> newRow = new HashMap<>(entry.getValue());
            localPheromoneMatrix.put(entry.getKey(), newRow);
        }
    }

    /**
     * 蚂蚁构建路径
     * @param user 用户信息
     * @param allTasks 所有任务
     * @param taskCacheMap 任务缓存
     * @param notIncreaseTaskIdSet 不增加质量的任务ID集合
     * @return 蚂蚁构建的路径
     */
    private List<TaskAlg> constructAntPath(UserAlg user, List<TaskAlg> allTasks,
                                           Map<Long, TaskAlg> taskCacheMap,
                                           Set<Long> notIncreaseTaskIdSet) {
        // 拷贝一份新的任务列表，避免修改原始数据
//        Map<Long,TaskAlg> candidateTaskMap = new HashMap<>();
        List<TaskAlg> candidateTasks = allTasks.stream().map(item -> {
            TaskAlg task = new TaskAlg();
            task.id = item.id;
            task.taskX = item.taskX;
            task.taskY = item.taskY;
            task.arrivalTime = item.arrivalTime;
            task.duration = item.duration;
            task.endTime = item.endTime;
            task.executionCost = item.executionCost;
            task.distanceCost = item.distanceCost;
            task.totalF = item.totalF;
            task.successProb = item.successProb;
            task.deltaF = item.deltaF;
            task.k = item.k;
            task.fCache = item.fCache;
//            candidateTaskMap.put(task.id, task);
            return task;
        }).collect(Collectors.toList());
        List<TaskAlg> antPath = new ArrayList<>();
        Set<Long> selectedTaskIdSet = new HashSet<>();

        double usedCost = 0;
        double tailX = user.startX;
        double tailY = user.startY;
        Long lastTaskId = null;

        while (true) {
            TaskAlg nextTask = selectNextTask(antPath, candidateTasks, selectedTaskIdSet,
                    notIncreaseTaskIdSet, tailX, tailY,
                    lastTaskId, user, usedCost,taskCacheMap);

            if (nextTask == null) {
                break;
            }

            // 计算添加该任务的成本
            double distanceCost;
            if (antPath.isEmpty()) {
                distanceCost = calculateDistance(user.startX, user.startY,
                        nextTask.taskX, nextTask.taskY);
            } else {
                distanceCost = calculateDistance(tailX, tailY,
                        nextTask.taskX, nextTask.taskY);
            }

            double totalCost = distanceCost + nextTask.executionCost;

            // 检查预算约束
            if (usedCost + totalCost + calculateFinalTrip(nextTask, user) > user.budget) {
                break;
            }

            // 放入缓存
            taskCacheMap.put(nextTask.id, nextTask);
            // 更新路径和状态
            antPath.add(nextTask);
            selectedTaskIdSet.add(nextTask.id);
            candidateTasks.remove(nextTask);
            usedCost += totalCost;
            tailX = nextTask.taskX;
            tailY = nextTask.taskY;
            // 更新局部信息素
            if (lastTaskId != null) {
                updateLocalPheromone(lastTaskId, nextTask.id);
            }
            lastTaskId = nextTask.id;
        }

        return antPath;
    }

    /**
     * 选择下一个任务
     */
    private TaskAlg selectNextTask(List<TaskAlg> currentPath, List<TaskAlg> candidateTasks,
                                   Set<Long> selectedTaskIdSet, Set<Long> notIncreaseTaskIdSet,
                                   double tailX, double tailY, Long lastTaskId,
                                   UserAlg user, double usedCost, Map<Long, TaskAlg> taskCacheMap) {
        if (candidateTasks.isEmpty()) {
            return null;
        }

        // 计算所有候选任务的偏好值
        Map<TaskAlg, Double> preferences = new HashMap<>();
        for (TaskAlg task : candidateTasks) {
            if (selectedTaskIdSet.contains(task.id) || notIncreaseTaskIdSet.contains(task.id)) {
                continue;
            }
            TaskAlg taskCache = taskCacheMap.get(task.id);
            if (taskCache != null && taskCache.k > 0){
                task.k = taskCache.k;
            }
            // 计算旅行成本
            double distanceCost;
            if (currentPath.isEmpty()) {
                distanceCost = calculateDistance(user.startX, user.startY,
                        task.taskX, task.taskY);
            } else {
                distanceCost = calculateDistance(tailX, tailY,
                        task.taskX, task.taskY);
            }

            double totalCost = distanceCost + task.executionCost;
            double weightCost = distanceCost + task.executionCost * alpha;
            // 检查预算约束
            if (usedCost + totalCost + calculateFinalTrip(task, user) > user.budget) {
                continue;
            }

            // 计算偏好值
            double preference = calculatePreference(lastTaskId, task, candidateTasks,
                    selectedTaskIdSet, notIncreaseTaskIdSet,
                    weightCost,currentPath);
            preferences.put(task, preference);
        }

        if (preferences.isEmpty()) {
            return null;
        }

        // 按偏好值排序
        List<Map.Entry<TaskAlg, Double>> sorted = new ArrayList<>(preferences.entrySet());
        sorted.sort((e1, e2) -> Double.compare(e2.getValue(), e1.getValue()));

        // 按概率选择任务
        if (Math.random() < bestSelectionProbability) {
            return sorted.get(0).getKey(); // 选择偏好值最高的任务
        } else {
            // 按偏好值比例选择其他任务
            double totalPreference = sorted.stream().mapToDouble(Map.Entry::getValue).sum();
            double randomValue = Math.random() * totalPreference;
            double cumulative = 0;

            for (Map.Entry<TaskAlg, Double> entry : sorted) {
                cumulative += entry.getValue();
                if (cumulative >= randomValue) {
                    return entry.getKey();
                }
            }

            return sorted.get(0).getKey(); // 默认返回第一个
        }
    }

    /**
     * 计算偏好值
     */
    private double calculatePreference(Long lastTaskId, TaskAlg candidate,
                                       List<TaskAlg> allTasks,
                                       Set<Long> selectedTaskIdSet,
                                       Set<Long> notIncreaseTaskIdSet,
                                       double weightCost, List<TaskAlg> currentPath) {
        // 计算信息素部分
        double pheromonePart = initialPheromone;
        if (lastTaskId != null) {
            Double pheromone = localPheromoneMatrix.get(lastTaskId).get(candidate.id);
            pheromonePart = Math.pow(pheromone != null ? pheromone : initialPheromone, pheromoneWeight);
        }

        // 使用QEA的calculateUtility方法计算启发式部分
        double heuristicUtility = super.calculateUtility(
                candidate,
                allTasks,
                weightCost,
                selectedTaskIdSet,
                notIncreaseTaskIdSet,
                candidate.calDeltaFCommon(), currentUser.budget, currentUser.usedCost);

        double heuristicPart = Math.pow(heuristicUtility, heuristicWeight);
        return pheromonePart * heuristicPart;
    }



    /**
     * 更新局部信息素
     */
    private void updateLocalPheromone(Long fromTaskId, Long toTaskId) {
        Map<Long, Double> row = localPheromoneMatrix.get(fromTaskId);
        if (row != null) {
            Double currentPheromone = row.get(toTaskId);
            if (currentPheromone != null) {
                // 信息素挥发
                row.put(toTaskId, (1 - evaporationRate) * currentPheromone + evaporationRate * initialPheromone);
            }
        }
    }

    /**
     * 更新全局信息素
     */
    private void updateGlobalPheromone(List<List<TaskAlg>> allAntPaths, List<Double> allAntPathUtilities) {
        // 找到最优的两条路径
        List<Integer> bestIndices = findBestPathIndices(allAntPathUtilities, 2);

        if (bestIndices.size() >= 1) {
            // 更新最优路径的信息素
            List<TaskAlg> bestPath = allAntPaths.get(bestIndices.get(0));
            updatePathPheromone(bestPath, initialPheromone);

            if (bestIndices.size() >= 2) {
                // 更新次优路径的信息素
                List<TaskAlg> secondBestPath = allAntPaths.get(bestIndices.get(1));
                updatePathPheromone(secondBestPath, initialPheromone / 2);
            }
        }
    }

    /**
     * 找到最优路径的索引
     */
    private List<Integer> findBestPathIndices(List<Double> utilities, int count) {
        List<Integer> indices = new ArrayList<>();
        PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparingDouble(i -> utilities.get(i)));

        for (int i = 0; i < utilities.size(); i++) {
            pq.offer(i);
            if (pq.size() > count) {
                pq.poll();
            }
        }

        while (!pq.isEmpty()) {
            indices.add(pq.poll());
        }

        Collections.reverse(indices);
        return indices;
    }

    /**
     * 更新路径信息素
     */
    private void updatePathPheromone(List<TaskAlg> path, double delta) {
        if (path.size() < 2) {
            return;
        }

        Long prevTaskId = path.get(0).id;
        for (int i = 1; i < path.size(); i++) {
            Long currentTaskId = path.get(i).id;

            Map<Long, Double> row = globalPheromoneMatrix.get(prevTaskId);
            if (row != null) {
                Double currentPheromone = row.get(currentTaskId);
                if (currentPheromone != null) {
                    row.put(currentTaskId, currentPheromone + delta);
                }
            }

            prevTaskId = currentTaskId;
        }
    }

    /**
     * 选择最优路径
     */
    private List<TaskAlg> selectBestPath(List<List<TaskAlg>> allAntPaths, List<Double> allAntPathUtilities) {
        if (allAntPaths.isEmpty()) {
            return new ArrayList<>();
        }

        int bestIndex = 0;
        double maxUtility = allAntPathUtilities.get(0);

        for (int i = 1; i < allAntPathUtilities.size(); i++) {
            if (allAntPathUtilities.get(i) > maxUtility) {
                maxUtility = allAntPathUtilities.get(i);
                bestIndex = i;
            }
        }

        return new ArrayList<>(allAntPaths.get(bestIndex));
    }

    /**
     * 更新任务执行次数和用户成本
     */
    private void updateTaskAndUserStatus(UserAlg user, List<TaskAlg> bestPath,
                                         Map<Long, TaskAlg> taskCacheMap,
                                         Set<Long> notIncreaseTaskIdSet) {
        double totalDistanceCost = 0;
        double totalExecutionCost = 0;

        // 计算路径总成本
        if (!bestPath.isEmpty()) {
            // 从起点到第一个任务的距离
            totalDistanceCost += calculateDistance(user.startX, user.startY,
                    bestPath.get(0).taskX, bestPath.get(0).taskY);

            // 任务间的距离
            for (int i = 1; i < bestPath.size(); i++) {
                totalDistanceCost += calculateDistance(bestPath.get(i-1).taskX, bestPath.get(i-1).taskY,
                        bestPath.get(i).taskX, bestPath.get(i).taskY);
            }

            // 从最后一个任务到终点的距离
            totalDistanceCost += calculateDistance(bestPath.get(bestPath.size()-1).taskX,
                    bestPath.get(bestPath.size()-1).taskY,
                    user.destX, user.destY);

            // 执行成本
            for (TaskAlg task : bestPath) {
                totalExecutionCost += task.executionCost;
            }
        }

        // 更新用户成本
        user.distanceCost = totalDistanceCost;
        user.executionCost = totalExecutionCost;
        user.usedCost = totalDistanceCost + totalExecutionCost;

        // 更新任务执行次数
        for (TaskAlg task : bestPath) {
            task.k++;
            TaskAlg taskCache = taskCacheMap.get(task.id);
            if (taskCache == null) {
                taskCache = task;
                taskCacheMap.put(task.id, taskCache);
            } else {
                taskCache.k++;
            }

            // 计算质量增量并更新总质量
            double deltaF = task.calDeltaFCommon();
            taskCache.totalF += deltaF;

            if (taskCache.totalF > 1) {
                log.info("任务[{}]的总质量大于1,totalF={},k={},successProb={}",
                        taskCache.id, taskCache.totalF, taskCache.k, taskCache.successProb);
                notIncreaseTaskIdSet.add(taskCache.id);
            }
        }
    }



    public static void main(String[] args) {
        Map<Long, Map<Long, Double>> localPheromoneMatrix = new HashMap<>();
        Map<Long, Map<Long, Double>> globalPheromoneMatrix = new HashMap<>();
        List<TaskAlg> tasks = new ArrayList<>();
        TaskAlg taskA = new TaskAlg();
        TaskAlg taskB = new TaskAlg();
        TaskAlg taskC = new TaskAlg();
        TaskAlg taskD = new TaskAlg();
        taskA.id = 1L;
        taskB.id = 2L;
        taskC.id = 3L;
        taskD.id = 4L;
        tasks.add(taskA);
        tasks.add(taskB);
        tasks.add(taskC);
        tasks.add(taskD);
        double initialPheromone = 0.1;

        // 初始化所有任务间的信息素
        for (TaskAlg task1 : tasks) {
            Map<Long, Double> row1 = new HashMap<>();
            Map<Long, Double> row2 = new HashMap<>();

            for (TaskAlg task2 : tasks) {
                if (!task1.equals(task2)) {
                    row1.put(task2.id, initialPheromone);
                    row2.put(task2.id, initialPheromone);
                }
            }
            globalPheromoneMatrix.put(task1.id, row1);
//            localPheromoneMatrix.put(task1.id, row2);
        }

        for (Map.Entry<Long, Map<Long, Double>> entry : globalPheromoneMatrix.entrySet()) {
            Map<Long, Double> newRow = new HashMap<>(entry.getValue());
            localPheromoneMatrix.put(entry.getKey(), newRow);
        }
        localPheromoneMatrix.get(1L).put(2L, 0.5); // 示例：设置特定任务间的信息素值
        System.out.println("localPheromoneMatrix:"+localPheromoneMatrix);
        System.out.println("globalPheromoneMatrix:"+globalPheromoneMatrix);


    }
}